package com.example.leetcode.graph.common.unionfind.twodimension;

/**
 * 二维并查集数据结构  quick-union的写法
 */
public class UnionFind {
    // 二维根数组大小
    int m, n;
    // 根数组，记录每个节点的父元素
    int[][][] root;
    // 记录每个根节点所代表集合中的元素数量
    int[][] size;

    public UnionFind(int m, int n) {
        this.m = m;
        this.n = n;
        this.root = new int[m][n][2];
        this.size = new int[m][n];
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                // 初始时每个元素的父结点就是自身
                root[i][j][0] = i;
                root[i][j][1] = j;
                // 初始时每个元素就是一个单独的集合
                size[i][j] = 1;
            }
        }
    }

    /**
     * 判断两个节点是否相连，即是否在同一个集合中
     *
     * @param i1
     * @param j1
     * @param i2
     * @param j2
     * @return
     */
    public boolean connected(int i1, int j1, int i2, int j2) {
        int[] root1 = find(i1, j1);
        int[] root2 = find(i2, j2);
        return root1[0] == root2[0] && root1[1] == root2[1];
    }

    /**
     * 查找某个节点的根节点
     *
     * @param i
     * @param j
     * @return
     */
    public int[] find(int i, int j) {
        int[] rootArr = root[i][j];
        int ri = rootArr[0], rj = rootArr[1];
        // 当节点的父结点是自身时，该节点是集合的根节点
        if (ri == i && rj == j) {
            return rootArr;
        }
        // 递归查找节点的父结点直到父结点是自身为止，即找到了根节点
        return find(ri, rj);
    }

    /**
     * 将两个集合合并成一个集合，小集合合并到大集合上面，将小集合挂到大集合上面
     *
     * @param i1
     * @param j1
     * @param i2
     * @param j2
     */
    public void union(int i1, int j1, int i2, int j2) {
        // 先查找当前两个元素的根节点
        int[] rootArr1 = find(i1, j1);
        int[] rootArr2 = find(i2, j2);
        int ri1 = rootArr1[0], rj1 = rootArr1[1];
        int ri2 = rootArr2[0], rj2 = rootArr2[1];
        // 将小集合的根节点挂在大集合的根节点下，合并两个节点所在的集合
        // 将节点数量少的树合并到节点数量多的树下
        // 也可以避免最终多次合并后形成一个链表（类似于按秩合并的优化思路，但是又不是按秩合并）
        // 但是效果是一样的，也是避免形成链表
        if (ri1 != ri2 || rj1 != rj2) {
            if (size[ri1][rj1] >= size[ri2][rj2]) {
                root[ri2][rj2][0] = ri1;
                root[ri2][rj2][1] = rj1;
                size[ri1][rj1] += size[ri2][rj2];  // size 每次累加后，就是树中的元素数量
            } else {
                root[ri1][rj1][0] = ri2;
                root[ri1][rj1][1] = rj2;
                size[ri2][rj2] += size[ri1][rj1];
            }
        }
    }
}
